package com.gdufe.firesafe.utils;

import io.github.briqt.spark4j.SparkClient;
import io.github.briqt.spark4j.constant.SparkApiVersion;
import io.github.briqt.spark4j.listener.SparkBaseListener;
import io.github.briqt.spark4j.model.SparkMessage;
import io.github.briqt.spark4j.model.SparkSyncChatResponse;
import io.github.briqt.spark4j.model.request.SparkRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * 各种向AI发送请求的方法
 */
@Component
@Slf4j
public class AIRequestUtils {

    @Resource
    private SparkClient sparkClient;

    /**
     * 默认采样温度
     */
    public static final Double DEFAULT_TEMPERATURE = 0.5;

    /**
     * 最大token数
     */
    @Value("xun-fei.maxToken")
    public static final Integer DEFAULT_MAX_TOKENS = null;

    /**
     * 向 AI 发送请求
     */
    public String sendMsgToXingHuo(boolean isNeedTemplate, String content) {
        if (isNeedTemplate) {
            // AI 生成问题的预设条件
            String predefinedInformation = "你是一个数据分析师和前端开发专家，接下来我会按照以下固定格式给你提供内容：\n" +
                    "分析需求：\n" +
                    "{数据分析的需求或者目标}\n" +
                    "原始数据：\n" +
                    "{csv格式的原始数据，用,作为分隔符}\n" +
                    "请根据这两部分内容，严格按照以下指定格式生成内容（此外不要输出任何多余的开头、结尾、注释）同时不要使用这个符号 '】'\n" +
                    "'【【【【【'\n" +
                    "{前端 Echarts V5 的 option 配置对象 JSON 代码, 不要生成任何多余的内容，比如注释和代码块标记}\n" +
                    "'【【【【【'\n" +
                    "{明确的数据分析结论、越详细越好，不要生成多余的注释} \n"
                    + "下面是一个具体的例子的模板："
                    + "'【【【【【'\n"
                    + "JSON格式代码"
                    + "'【【【【【'\n" +
                    "结论：";
            content = predefinedInformation + "\n" + content;
        }
        List<SparkMessage> messages = new ArrayList<>();
        messages.add(SparkMessage.userContent(content));
        // 构造请求
        SparkRequest sparkRequest = SparkRequest.builder()
                // 消息列表
                .messages(messages)
                // 模型回答的tokens的最大长度,非必传,取值为[1,4096],默认为2048
                .maxTokens(500)
                // 核采样阈值。用于决定结果随机性,取值越高随机性越强即相同的问题得到的不同答案的可能性越高 非必传,取值为[0,1],默认为0.5
                .temperature(0.2)
                // 指定请求版本，默认使用最新2.0版本
                .apiVersion(SparkApiVersion.V3_5)
                .build();
        // 同步调用
        SparkSyncChatResponse chatResponse = sparkClient.chatSync(sparkRequest);
        String responseContent = chatResponse.getContent();
        log.info("星火 AI 返回的结果 {}", responseContent);
        return responseContent;
    }

    /**
     * 传递用户消息的异步流式请求，默认采样温度
     * @param systemMessage
     */
    public void streamRequest(String systemMessage, SparkBaseListener listener) {
        List<SparkMessage> messages = new ArrayList<>();
        messages.add(SparkMessage.userContent(systemMessage));
        log.debug("系统提示词：\n{}", systemMessage);
        streamRequest(messages, DEFAULT_TEMPERATURE, listener);
    }

    /**
     * 传递系统消息和用户消息的异步流式请求，默认采样温度
     * @param systemMessage
     * @param userMessage
     */
    public void streamRequest(String systemMessage, String userMessage, SparkBaseListener listener) {
        List<SparkMessage> messages = new ArrayList<>();
        messages.add(SparkMessage.userContent(userMessage));
        messages.add(SparkMessage.systemContent(systemMessage));
        log.debug("系统提示词：\n{}", systemMessage);
        log.debug("用户提示词：\n{}", userMessage);
        streamRequest(messages, DEFAULT_TEMPERATURE, listener);
    }

    /**
     * 通用异步流式请求
     * @param messages
     * @param temperature
     */
    public void streamRequest(List<SparkMessage> messages, Double temperature, SparkBaseListener listener) {
        SparkRequest sparkRequest = buildRequest(messages, temperature);
        sparkClient.chatStream(sparkRequest, listener);
    }

    /**
     * 传递系统消息和用户消息的同步请求，默认采样温度
     * @param systemMessage
     * @param userMessage
     * @return
     */
    public String syncRequest(String systemMessage, String userMessage){
        List<SparkMessage> messages = new ArrayList<>();
        messages.add(SparkMessage.userContent(userMessage));
        messages.add(SparkMessage.systemContent(systemMessage));
        log.debug("系统提示词：\n{}", systemMessage);
        log.debug("用户提示词：\n{}", userMessage);
        return syncRequest(messages, DEFAULT_TEMPERATURE);
    }

    /**
     * 通用同步请求
     *
     * @param messages
     * @param temperature
     * @return
     */
    public String syncRequest(List<SparkMessage> messages, Double temperature) {
        // 构造请求
        SparkRequest sparkRequest = buildRequest(messages, temperature);
        String responseOfAI = sparkClient.chatSync(sparkRequest).getContent();
        log.info("AI返回结果：\n{}", responseOfAI);
        return responseOfAI;
    }

    /**
     * 构造请求
     * @param messages
     * @param temperature
     * @return
     */
    private static SparkRequest buildRequest(List<SparkMessage> messages, Double temperature) {
        return SparkRequest.builder()
                // 消息列表
                .messages(messages)
                // 模型回答的tokens的最大长度,非必传,取值为[1,4096],默认为2048
                .maxTokens(DEFAULT_MAX_TOKENS)
                // 核采样阈值。用于决定结果随机性,取值越高随机性越强即相同的问题得到的不同答案的可能性越高 非必传,取值为[0,1],默认为0.5
                .temperature(temperature)
                // 指定请求版本，默认使用最新2.0版本
                .apiVersion(SparkApiVersion.V4_0)
                .build();
    }

}
